import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler


def pca():
    df = pd.read_csv('iris.csv')
    x = df.iloc[:, 1:5].values
    y = df.iloc[:, 5].values

    # 数据标准化
    scaler = StandardScaler()
    x_std = scaler.fit_transform(x)

    # 手动实现PCA
    # 求协方差矩阵
    cov_mat = np.cov(x_std.T)

    # 计算特征值和特征向量
    eig_vals, eig_vecs = np.linalg.eig(cov_mat)

    # 将特征值和对应的特征向量配对
    eig_pairs = [(np.abs(eig_vals[i]), eig_vecs[:, i]) for i in range(len(eig_vals))]

    # 按特征值大小降序排序
    eig_pairs.sort(key=lambda x: x[0], reverse=True)

    # 选择前两个主成分
    matrix_w = np.hstack((eig_pairs[0][1].reshape(4, 1),
                          eig_pairs[1][1].reshape(4, 1)))

    # 将数据投影到主成分上
    Y = x_std.dot(matrix_w)

    # 创建对比图像
    plt.figure(figsize=(14, 6))

    # 子图1：原始数据的前两个特征
    plt.subplot(1, 2, 1)
    colors = {'Iris-setosa': 'blue', 'Iris-versicolor': 'red', 'Iris-virginica': 'green'}
    for lab, col in colors.items():
        plt.scatter(x[y == lab, 0],
                    x[y == lab, 1],
                    label=lab,
                    c=col,
                    alpha=0.7,
                    s=60,
                    edgecolors='black',
                    linewidth=0.5)
    plt.xlabel('Sepal Length (cm)')
    plt.ylabel('Sepal Width (cm)')
    plt.title('Original Data - First Two Features')
    plt.legend(loc='best')
    plt.grid(True, alpha=0.3)

    # 子图2：PCA降维后的数据
    plt.subplot(1, 2, 2)
    for lab, col in colors.items():
        plt.scatter(Y[y == lab, 0],
                    Y[y == lab, 1],
                    label=lab,
                    c=col,
                    alpha=0.7,
                    s=60,
                    edgecolors='black',
                    linewidth=0.5)
    plt.xlabel('Principal Component 1')
    plt.ylabel('Principal Component 2')
    plt.title('After PCA')
    plt.legend(loc='best')
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('pca_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()

    return Y


if __name__ == '__main__':
    pca_result = pca()